package com.shiguiwu.langchain.config;

import com.google.common.collect.Lists;
import com.shiguiwu.langchain.service.ChatService;
import dev.langchain4j.data.document.DefaultDocument;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @description: 聊天配置
 * @author: stone
 * @date: Created by 2025/6/19 17:16
 * @version: 1.0.0
 * @pakeage: com.shiguiwu.langchain.config
 */

@Configuration
public class ChatConfig {


    @Autowired
    private EmbeddingModel embeddingModel;
    /**
     * 编程式注入
     * @param model
     * @return
     */
    @Bean
    public ChatService chatService(OpenAiChatModel model) {
        return AiServices.builder(ChatService.class)
                .chatModel(model)//设置对话时使用的模型对象
                .build();
    }

    /**
     * 会话记忆
     * @return
     */
    @Bean
    public ChatMemory chatMemory() {
        return MessageWindowChatMemory.builder()
                .maxMessages(20)//最大保存的会话记录数量
                .build();
    }

    public ChatMemory chatMemory(Object memoryId) {
        return MessageWindowChatMemory.builder()
                .maxMessages(20)//最大保存的会话记录数量
                .id(memoryId)
                .build();
    }

    public ChatMemory chatMemory(Object memoryId, ChatMemoryStore redisChatMemoryStore) {

        return MessageWindowChatMemory.builder()
                .maxMessages(20)//最大保存的会话记录数量
                .id(memoryId)
                .chatMemoryStore(redisChatMemoryStore)
                .build();
    }

    @Bean
    public  ChatMemoryProvider chatMemoryProvider() {
        return this::chatMemory;
    }

    @Autowired
    private ChatMemoryStore redisChatMemoryStore;


    @Bean
    public  ChatMemoryProvider chatRedisMemoryProvider() {
        return (memoryId) -> this.chatMemory(memoryId, redisChatMemoryStore);
    }




    //@Bean
    public EmbeddingStore store1(){
        //1.加载文档进内存
        List<Document> documents = ClassPathDocumentLoader.loadDocuments("content", new ApachePdfBoxDocumentParser());
        //2.构建向量数据库操作对象  操作的是内存版本的向量数据库
        //InMemoryEmbeddingStore store = new InMemoryEmbeddingStore<>();
        QdrantEmbeddingStore docsVectors = QdrantEmbeddingStore.builder()
                .host("47.120.77.92")
                .port(6334)
                .apiKey("12345678")
                .collectionName("docs_vectors")

                .build();
        //构建文档分割器对象
        DocumentSplitter ds = DocumentSplitters.recursive(500,100);
        //3.构建一个EmbeddingStoreIngestor对象,完成文本数据切割,向量化, 存储
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .embeddingStore(docsVectors)
                .documentSplitter(ds)
                .embeddingModel(embeddingModel)

                .build();
        // 5. 分批处理文档（关键修改）
        for (Document doc : documents) {
            // 单独处理每个文档
            //List<Document> singleDocList = Collections.singletonList(doc);

            // 获取分割后的文本块
            List<TextSegment> segments = ds.split(doc);

            // 分批处理文本块（每批最多10个）
            List<List<TextSegment>> batches = Lists.partition(segments, 10);

            for (List<TextSegment> batch : batches) {
                // 转换为Document列表
                List<Document> batchDocs = batch.stream()
                        .map(seg -> new DefaultDocument(seg.text(), seg.metadata()))
                        .collect(Collectors.toList());

                // 处理当前批次
                ingestor.ingest(batchDocs);
            }
        }

        return docsVectors;
    }

    @Bean
    public EmbeddingStore store(){
        return QdrantEmbeddingStore.builder()
                .host("47.120.77.92")
                .port(6334)
                .apiKey("12345678")
                .collectionName("docs_vectors")
                .build();
    }

    @Bean
    public ContentRetriever contentRetriever(EmbeddingStore store) {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(store)//设置向量数据库操作对象
                .minScore(0.6)//设置最小分数
                .maxResults(3)//设置最大片段数量
                .embeddingModel(embeddingModel)
                .build();
    }

}